#!/usr/bin/env python3
# -*- coding:utf-8 -*-

import os
import sys
import time
import getopt
import socket
import subprocess


def execcmd(cmd):
    print("exec cmd:", cmd)
    state, output = subprocess.getstatusoutput(cmd)
    
    if state != 0:
        print('命令执行遇到错误，错误码：', state)
        print('错误内容：', output)
        Report.print_and_exit(1)
    
    return output


class Report(object):
    step = 0
    stages = ['初次升级', '初次回滚', '再次升级', '升级提交']
    results = ['待检', '待检', '待检', '待检']
    reports = ['', '', '', '']
    
    @staticmethod
    def record(res, report):
        i = Report.step
        Report.step += 1
        
        archive = f'{ProjCxt.workspace}/{Report.stages[i]}-report.md'
        cmd = f'cp {report} {archive}'
        execcmd(cmd)
        
        Report.results[i] = res
        Report.reports[i] = archive

    @staticmethod
    def print_and_exit(exit_code):
        print('升级四步校验结果：')
        for i, step in enumerate(Report.stages):
            print(f'    {step}： {Report.results[i]}  {Report.reports[i]}')
        exit(exit_code)
        

class Xml(object):
    def __init__(self, username, ip, port, file):
        self.file = file
        self.username = username
        self.ip = ip
        self.port = port
        self.hostname= socket.gethostname()
        content = f"""<?xml version="1.0" encoding="UTF-8"?>
<ROOT>
  <CLUSTER>
    <PARAM name="clusterName" value="gauss_upgrade"/>
    <PARAM name="nodeNames" value="{self.hostname}"/>
    <PARAM name="gaussdbAppPath" value="/home/{username}/cluster/app"/>
    <PARAM name="gaussdbLogPath" value="/home/{username}/cluster/gaussdb_log" />
    <PARAM name="tmpMppdbPath" value="/home/{username}/cluster/tmp"/>
    <PARAM name="gaussdbToolPath" value="/home/{username}/cluster/tool"/>
    <PARAM name="corePath" value="/home/{username}/core"/>
    <PARAM name="backIp1s" value="{ip}"/>
  </CLUSTER>
  <DEVICELIST>
    <DEVICE sn="{self.hostname}">
        <PARAM name="name" value="{self.hostname}"/>
        <PARAM name="backIp1" value="{self.ip}"/>
        <PARAM name="sshIp1" value="{self.ip}"/>
        <PARAM name="azName" value="AZ1"/>
        <PARAM name="azPriority" value="1"/>
        <!-- dn -->
        <PARAM name="dataNum" value="1"/>
        <PARAM name="dataPortBase" value="{port}"/>
        <PARAM name="dataNode1" value="/home/{username}/cluster/dn1"/>
    </DEVICE>
  </DEVICELIST>
</ROOT>
"""
        with open(file, 'w') as f:
            f.write(content)


class UserCxt(object):
    username = None
    password = 'SFRXjwdd@250'
    fi_old = None
    fi_old_env = None
    fi_old_port = None
    fi_new = None
    fi_new_env = None
    fi_new_port = None
    
    
    def init(username):
        UserCxt.username = username
        
        UserCxt.fi_old = f'/home/{username}/fastinstallold'
        UserCxt.fi_old_env = f'{UserCxt.fi_old}/env.source'
        UserCxt.fi_old_port = ProjCxt.port + 100
        UserCxt.fi_new = f'/home/{username}/fastinstallnew'
        UserCxt.fi_new_env = f'{UserCxt.fi_new}/env.source'
        UserCxt.fi_new_port = ProjCxt.port + 200
        
        try:
            cmd = f"id {username}"
            subprocess.run(cmd, shell=True, check=True)
            print(f'用户{username}早已存在。可以使用：userdel {username} 删除。')
            exit(1)
        except subprocess.CalledProcessError:
            print(f'创建用户{username}。')
        execcmd('useradd ' + username)
        execcmd(f'echo "{username}:{UserCxt.password}" | chpasswd')


class ProjCxt(object):
    useage = """
python3 upgrade-fastcheck.py -o -n -u -i -p [-d]
    -o --old oldpackage
    -n --new newpackage
    -u --user user
    -i --ip ip
    -p --port port
    -d  --dir workspace
    """
    
    # param
    old_pkg = None
    new_pkg = None
    user = None
    ip = None
    port = None
    workspace = None

    xml = None
    pkgspace = None
    pkgspace_old = None
    pkgspace_new = None
    
    @staticmethod
    def init():
        ProjCxt._parse_param()
        
        ProjCxt.pkgspace = os.path.join(ProjCxt.workspace, 'packages')
        ProjCxt.pkgspace_old = os.path.join(ProjCxt.pkgspace, 'oldpkg')
        ProjCxt.pkgspace_new = os.path.join(ProjCxt.pkgspace, 'newpkg')
        for directory in [ProjCxt.workspace, ProjCxt.pkgspace, ProjCxt.pkgspace_new, ProjCxt.pkgspace_old]:
            if not os.access(directory, os.F_OK):
                os.mkdir(directory, 0o755)
        
        xmlpath = os.path.join(ProjCxt.workspace, 'upgrade-fastcheck.xml')
        ProjCxt.xml = Xml(ProjCxt.user, ProjCxt.ip, ProjCxt.port, xmlpath)
        
        UserCxt.init(ProjCxt.user)
        untar_pkg(ProjCxt.old_pkg, ProjCxt.pkgspace_old)
        untar_pkg(ProjCxt.new_pkg, ProjCxt.pkgspace_new)
        
        
    @staticmethod
    def destroy():
        """userdel carrot2"""
        cmd = f"su - {ProjCxt.user} -c 'killall gaussdb' "
        state, output = subprocess.getstatusoutput(cmd)
        
        cmd = f'rm /home/{ProjCxt.user} -fr'
        state, output = subprocess.getstatusoutput(cmd)
        
        
    @staticmethod
    def _parse_param():
        argv = sys.argv[1:]
        short_options = 'o:n:u:d:i:p:'
        long_options = ['old=', 'new=', 'user=', 'dir=', 'ip=', 'port=']
        opts, args = getopt.getopt(argv, short_options, long_options)
        if len(args) > 1:
            print('unknown param ' + args[0])
            exit(1)
        for opt, arg in opts:
            if opt in ['-o', '--old']:
                ProjCxt.old_pkg = arg
            elif opt in ['-n', '--new']:
                ProjCxt.new_pkg = arg
            elif opt in ['-x', '--xml']:
                ProjCxt.xml = arg
            elif opt in ['-u', '--user']:
                ProjCxt.user = arg
            elif opt in ['-d', '--dir']:
                ProjCxt.workspace = arg
            elif opt in ['-i', '--ip']:
                ProjCxt.ip = arg
            elif opt in ['-p', '--port']:
                ProjCxt.port = int(arg)
            else:
                assert False
        if ProjCxt.old_pkg is None or ProjCxt.new_pkg is None or ProjCxt.ip is None or ProjCxt.user is None or ProjCxt.port is None:
            print('缺少参数。\n', ProjCxt.useage)
            exit(1)
        if ProjCxt.workspace is None:
            ProjCxt.workspace = os.path.join(os.path.dirname(sys.argv[0]), 'upgrade-fastcheck-workspace')
        ProjCxt.workspace = os.path.abspath(ProjCxt.workspace)



def untar_pkg(pkg, directory):
    print(f'解压{pkg}到{directory}')
    output = execcmd(f'tar -zxvf {pkg} -C {directory}')
    
    count = 0
    subpkgs = output.splitlines()
    for subpkg in subpkgs:
        if subpkg.endswith('tar.bz2'):
            execcmd(f'cd {directory} && tar -jxvf {subpkg}')
            count += 1
        elif subpkg.endswith('om.tar.gz'):
            execcmd(f'cd {directory} && tar -zxvf {subpkg}')
            count += 1
            
    assert count == 2, '少包了。'


def fastinstall(package, user, path, envpath, port):
    """
    为某个用户，快速通过配置环境变量安装一个数据库
    """
    print(f'简易安装一个数据库。{package}  -> {path}')
    execcmd(f'mkdir {path}')
    
    # 解压安装包
    untar_pkg(package, path)
    
    # 创建环境变量
    env = f"""
    export GAUSSHOME={path}
    export GAUSSPORT={port}
    export PATH=$GAUSSHOME/bin:$GAUSSHOME/script:$PATH
    export LD_LIBRARY_PATH=$GAUSSHOME/lib:$LD_LIBRARY_PATH
    """
    with open(envpath, 'w') as f:
        f.write(env)
    execcmd(f'chown {user}:{user} /home/{user}/* -R')
    
    dnpath = f'{path}/sgnode'
    initdb_cmd = f'source {envpath} && gs_initdb -D {dnpath} -w BJLBJL@666 --nodename="sgnode"'
    cmd = f"su - {ProjCxt.user} -c '{initdb_cmd}' "
    execcmd(cmd)
    
    setport_cmd = f'''source {envpath} && gs_guc set -D {dnpath} -c "port='{port}'" '''
    cmd = f"su - {ProjCxt.user} -c '{setport_cmd}' "
    execcmd(cmd)
    
    rundb_cmd = f'source {envpath} && gs_ctl start -D {dnpath} '
    cmd = f"su - {ProjCxt.user} -c '{rundb_cmd}' "
    execcmd(cmd)
    
    
def pre_install(package):
    """
    预安装
    """
    script_path = os.path.join(package, 'script')
    cmd = f'cd {script_path} && ./gs_preinstall -U {ProjCxt.user} -G {ProjCxt.user} -X {ProjCxt.xml.file}'
    print('预安装', cmd)
    
    process = subprocess.Popen(cmd, stdin=subprocess.PIPE, stdout=subprocess.PIPE, shell=True)
    stdout, stderr = process.communicate(input=b'no\n')
    if process.returncode != 0:
        print(stdout.decode('utf-8'))
        print('命令执行遇到错误。', stderr.decode('utf-8'))
        Report.print_and_exit(1)
        exit(1)
    
    
def install():
    """
    安装
    """
    cmd = f"su - {ProjCxt.user} -c 'gs_install -X {ProjCxt.xml.file}'"
    process = subprocess.Popen(cmd, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True)
    print('安装', cmd)

    password = b"BJLbjl@123\n"
    process.stdin.write(password)
    process.stdin.flush()
    process.stdin.write(password)
    process.stdin.flush()

    stdout, stderr = process.communicate()

    if process.returncode != 0:
        print(stdout.decode('utf-8'))
        print('命令执行遇到错误。', stderr.decode('utf-8'))
        Report.print_and_exit(1)
        exit(1)


def upgrade_upgrade():
    """
    升级
    """
    cmd = f"su - {ProjCxt.user} -c 'gs_upgradectl -t auto-upgrade -X {ProjCxt.xml.file} --grey'"
    print('升级操作: ', cmd)
    execcmd(cmd)

def upgrade_rollback():
    """
    升级回滚
    """
    cmd = f"su - {ProjCxt.user} -c 'gs_upgradectl -t auto-rollback -X {ProjCxt.xml.file}'"
    print('升级回滚: ', cmd)
    state, output = subprocess.getstatusoutput(cmd)
    # todo  正常结束，但退出是2？暂时用日志关键字判断。
    if output.find('Rollback succeeded') < 0:
        print(output)
        print('命令执行遇到错误：', state)
        Report.print_and_exit(1)
        exit(1)

def upgrade_commit():
    """
    升级提交
    """
    print('升级提交')
    cmd = f"su - {ProjCxt.user} -c 'gs_upgradectl -t commit-upgrade -X {ProjCxt.xml.file}'"
    execcmd(cmd)


def export(port, envfile=None):
    """
    导出vmap
    """
    print('VMAP导出')
    sourcecmd = f'source {envfile} &&' if envfile is not None  else ' '
    cmd = f"su - {ProjCxt.user} -c '{sourcecmd} gs_upgradechk export -p {port}'"
    output = execcmd(cmd)
    vmap = output.splitlines()[-1].split(':')[1]
    print(vmap)
    return vmap


def verify(port, vmap):
    """
    校验
    """
    print('执行校验')
    cmd = f"su - {ProjCxt.user} -c 'gs_upgradechk verify -p {port} -v {vmap}'"
    output = execcmd(cmd)
    report = output.splitlines()[-1].split("：")[-1]
    
    with open(report, 'r') as f:
        lines = f.reallines()
    for line in lines:
        if not line.startswith('**校验结果**: '):
            continue
        res = line.split(': ')[1]
        
        Report.record(res, report)
        
        if res != '成功':
            Report.print_and_exit(1)
        return


if __name__ == '__main__':
    ProjCxt.init()
    
    # 生成两个vmap
    fastinstall(ProjCxt.old_pkg, ProjCxt.user, UserCxt.fi_old, UserCxt.fi_old_env, UserCxt.fi_old_port)
    old_vmap = export(UserCxt.fi_old_port, UserCxt.fi_old_env)
    fastinstall(ProjCxt.new_pkg, ProjCxt.user, UserCxt.fi_new, UserCxt.fi_new_env, UserCxt.fi_new_port)
    new_vmap = export(UserCxt.fi_new_port, UserCxt.fi_new_env)
    
    # 安装老版本
    pre_install(ProjCxt.pkgspace_old)
    install()
    
    # 1、升级到新版本并校验
    pre_install(ProjCxt.pkgspace_new)
    upgrade_upgrade()
    verify(ProjCxt.port, new_vmap)
    
    # 2、回滚到老版本并校验
    upgrade_rollback()
    verify(ProjCxt.port, old_vmap)
    
    # 3、再次升级到新版本并校验
    pre_install(ProjCxt.pkgspace_new)
    upgrade_upgrade()
    verify(ProjCxt.port, new_vmap)
    
    # 4、提交升级并校验
    upgrade_commit()
    verify(ProjCxt.port, new_vmap)

    # 清理
    ProjCxt.destroy()